package com.pmt.feedplayer.core

import android.media.MediaExtractor
import android.media.MediaFormat
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.flow
import kotlinx.coroutines.withContext
import java.io.IOException
import java.nio.ByteBuffer

/**
 * Android平台音频提取器实现
 */
class AndroidAudioExtractor : AudioExtractor {
    private var mediaExtractor: MediaExtractor? = null
    private var bufferSize = 1000 // 默认1000ms缓冲区
    private var isRunning = false
    
    override suspend fun extractAudio(videoUrl: String): Flow<ByteArray> = flow {
        isRunning = true
        try {
            withContext(Dispatchers.IO) {
                mediaExtractor = MediaExtractor()
                mediaExtractor?.setDataSource(videoUrl)
                
                // 查找音频轨道
                val audioTrackIndex = findAudioTrack(mediaExtractor!!)
                if (audioTrackIndex < 0) {
                    throw IOException("No audio track found in the video")
                }
                
                // 选择音频轨道
                mediaExtractor?.selectTrack(audioTrackIndex)
                
                // 获取音频格式
                val audioFormat = mediaExtractor?.getTrackFormat(audioTrackIndex)
                val sampleRate = audioFormat?.getInteger(MediaFormat.KEY_SAMPLE_RATE) ?: 44100
                val channels = audioFormat?.getInteger(MediaFormat.KEY_CHANNEL_COUNT) ?: 1
                
                // 计算每个缓冲区的样本数
                val samplesPerBuffer = (sampleRate * bufferSize / 1000) * channels
                val buffer = ByteArray(samplesPerBuffer)
                
                // 提取音频数据
                while (isRunning) {
                    // 使用ByteBuffer而不是ByteArray
                    val byteBuffer = ByteBuffer.allocate(samplesPerBuffer)
                    val sampleSize = mediaExtractor?.readSampleData(byteBuffer, 0) ?: -1
                    if (sampleSize <= 0) {
                        // 读取完毕
                        break
                    }
                    
                    // 将ByteBuffer转换回ByteArray
                    byteBuffer.flip()
                    val audioData = ByteArray(sampleSize)
                    byteBuffer.get(audioData)
                    
                    // 发送音频数据块
                    emit(audioData)
                    
                    // 前进到下一个样本
                    mediaExtractor?.advance()
                    
                    // 短暂延迟控制处理速度
                    Thread.sleep(10) // 调整延迟以控制提取速度
                }
            }
        } catch (e: Exception) {
            if (isRunning) { // 只有在运行中发生的异常才重新抛出
                throw e
            }
        } finally {
            release()
        }
    }
    
    override fun stop() {
        isRunning = false
        release()
    }
    
    override fun setBufferSize(sizeInMs: Int) {
        this.bufferSize = sizeInMs
    }
    
    private fun findAudioTrack(extractor: MediaExtractor): Int {
        for (i in 0 until extractor.trackCount) {
            val format = extractor.getTrackFormat(i)
            val mime = format.getString(MediaFormat.KEY_MIME)
            if (mime?.startsWith("audio/") == true) {
                return i
            }
        }
        return -1
    }
    
    private fun release() {
        mediaExtractor?.release()
        mediaExtractor = null
    }
}